package com.yunkeji.api.test;

import java.util.*;

public class AhoCorasick {
    private final Node root;

    public AhoCorasick() {
        root = new Node();
    }

    // 添加关键词到Trie树
    public void addKeyword(String keyword) {
        Node currentNode = root;
        for (char ch : keyword.toCharArray()) {
            currentNode = currentNode.getChildren().computeIfAbsent(ch, c -> new Node());
        }
        currentNode.setEndOfWord(true);
    }

    // 构建失配指针（Failure Links）
    public void buildFailureLinks() {
        Queue<Node> queue = new LinkedList<>();
        root.setFailureLink(root);

        for (Node child : root.getChildren().values()) {
            child.setFailureLink(root);
            queue.add(child);
        }

        while (!queue.isEmpty()) {
            Node currentNode = queue.poll();

            for (Map.Entry<Character, Node> entry : currentNode.getChildren().entrySet()) {
                char ch = entry.getKey();
                Node child = entry.getValue();
                Node failure = currentNode.getFailureLink();

                while (failure != root && !failure.getChildren().containsKey(ch)) {
                    failure = failure.getFailureLink();
                }

                if (failure.getChildren().containsKey(ch)) {
                    child.setFailureLink(failure.getChildren().get(ch));
                } else {
                    child.setFailureLink(root);
                }

                if (child.getFailureLink().isEndOfWord()) {
                    child.setEndOfWord(true);
                }

                queue.add(child);
            }
        }
    }

    // 执行关键词匹配
    public List<String> search(String text) {
        List<String> results = new ArrayList<>();
        Node currentNode = root;

        for (int i = 0; i < text.length(); i++) {
            char ch = text.charAt(i);

            while (currentNode != root && !currentNode.getChildren().containsKey(ch)) {
                currentNode = currentNode.getFailureLink();
            }

            if (currentNode.getChildren().containsKey(ch)) {
                currentNode = currentNode.getChildren().get(ch);
            } else {
                currentNode = root;
            }

            if (currentNode.isEndOfWord()) {
                results.add("Pattern found at index " + (i + 1));
            }
        }

        return results;
    }

    // Trie树节点定义
    private static class Node {
        private final Map<Character, Node> children = new HashMap<>();
        private Node failureLink;
        private boolean endOfWord;

        public Map<Character, Node> getChildren() {
            return children;
        }

        public Node getFailureLink() {
            return failureLink;
        }

        public void setFailureLink(Node failureLink) {
            this.failureLink = failureLink;
        }

        public boolean isEndOfWord() {
            return endOfWord;
        }

        public void setEndOfWord(boolean endOfWord) {
            this.endOfWord = endOfWord;
        }
    }

    public static void main(String[] args) {
        AhoCorasick ac = new AhoCorasick();
        ac.addKeyword("he");
        ac.addKeyword("she");
        ac.addKeyword("his");
        ac.addKeyword("hers");

        ac.buildFailureLinks();

        String text = "ushers";
        List<String> results = ac.search(text);

        for (String result : results) {
            System.out.println(result);
        }
    }
}
